import numpy as np
import matplotlib.pyplot as plt
from scipy.constants import g
from scipy.integrate import solve_ivp
from matplotlib.widgets import Slider, Button
import matplotlib as mpl
from matplotlib.patches import FancyBboxPatch
import matplotlib.patheffects as path_effects
from matplotlib.lines import Line2D
import warnings
import matplotlib.cm as cm
from scipy.interpolate import make_interp_spline
from matplotlib.ticker import MaxNLocator
from matplotlib.gridspec import GridSpec
import matplotlib.ticker as ticker

# 过滤solve_ivp的警告
warnings.filterwarnings('ignore', category=UserWarning, message='The following arguments have no effect')

# 设置微软雅黑字体
plt.rcParams['font.sans-serif'] = ['.HeitiUI TC','Microsoft YaHei', 'SimHei', 'WenQuanYi Micro Hei']
plt.rcParams['axes.unicode_minus'] = False

# === 初始水火箭参数 ===
initial_params = {
    'dry_mass': 0.070,          # kg
    'water_fraction': 0.40,     # 瓶容积比例
    'initial_pressure': 3.5,   # atm (内部使用)
    'launch_angle': 50,         # 度
    'bottle_volume': 0.54,      # 升
    'bottle_diameter': 0.10,   # 米
    'nozzle_diameter': 0.013,  # 米
    'drag_coefficient': 0.5    # 无量纲
}

# PSI到ATM的转换系数
PSI_TO_ATM = 14.6959

# 当前参数
params = initial_params.copy()

# 物理常量
rho_water = 1000       # kg/m³
p0 = 101325            # Pa
rho_air = 1.225        # kg/m³
gamma = 1.4            # 绝热指数

# === 物理模型 ===
def simulate_rocket(params):
    # 转换为SI单位
    V_bottle = params['bottle_volume'] * 0.001  # m³
    S2 = np.pi * (params['bottle_diameter']/2)**2  # m²
    l0 = V_bottle / S2  # m
    S1 = np.pi * (params['nozzle_diameter']/2)**2  # m²
    pi = params['initial_pressure'] * p0  # Pa
    global_theta = np.radians(params['launch_angle'])
    
    # 初始条件
    initial_l = params['water_fraction'] * l0  # m
    Vi = (l0 - initial_l) * S2  # m³
    
    # 安全检查
    if initial_l <= 0 or initial_l >= l0:
        initial_l = 0.4 * l0
        Vi = (l0 - initial_l) * S2
    
    # 存储阶段信息
    phases = []
    
    def model(t, y):
        l, vx, vy, x, y_pos = y
        
        # 火箭总质量
        water_mass = rho_water * S2 * l
        m = params['dry_mass'] + water_mass
        
        # 空气体积和压力
        V_air = max(1e-6, (l0 - l) * S2)
        p = pi * (Vi/V_air)**gamma
        
        # 水流条件
        water_flowing = (l > 1e-6) and (y_pos > 0) and (p > p0)
        
        # 水流计算
        if water_flowing:
            pressure_diff = max(0, p - p0)
            v2 = np.sqrt(2 * pressure_diff / rho_water)
            dm_dt = -rho_water * S1 * v2
            thrust = -dm_dt * v2
        else:
            v2 = 0
            dm_dt = 0
            thrust = 0
        
        # 气动阻力
        v_speed = np.sqrt(vx**2 + vy**2)
        if v_speed > 0.01:
            eta = 0.5 * params['drag_coefficient'] * rho_air * S2
            Fr = eta * v_speed**2
            Fr_x = -Fr * (vx/v_speed)
            Fr_y = -Fr * (vy/v_speed)
        else:
            Fr_x, Fr_y = 0, 0
        
        # 推力方向 - 始终沿速度方向
        if thrust > 0:
            if v_speed > 0.01:
                thrust_x = thrust * (vx/v_speed)
                thrust_y = thrust * (vy/v_speed)
            else:
                # 初始时刻使用发射角度
                thrust_x = thrust * np.cos(global_theta)
                thrust_y = thrust * np.sin(global_theta)
        else:
            thrust_x, thrust_y = 0, 0
        
        # 加速度方程
        dvx_dt = (thrust_x + Fr_x) / m
        dvy_dt = (thrust_y - m*g + Fr_y) / m
        
        # 水位变化
        dl_dt = dm_dt/(rho_water*S2) if water_flowing else 0
        
        # 地面碰撞处理
        if y_pos <= 0 and vy < 0:
            dvx_dt = 0
            dvy_dt = 0
            vx = 0
            vy = 0
        
        return [dl_dt, dvx_dt, dvy_dt, vx, vy]
    
    # 初始状态 - 将初始高度改为0.15米 (15cm)
    initial_state = [
        initial_l,
        0,  # vx
        0,  # vy
        0,  # x
        0.15  # y (15cm)
    ]
    
    # 时间设置
    t_max = 15.0
    t_eval = np.linspace(0, t_max, 10000)
    
    try:
        sol = solve_ivp(model, [0, t_max], initial_state, t_eval=t_eval,
                        method='RK45', rtol=1e-6, atol=1e-6)
        t = sol.t
        # 确保解包正确
        if sol.y.shape[0] >= 5:
            l, vx, vy, x, y = sol.y[:5]
        else:
            # 如果解包失败，使用默认值
            l = np.zeros_like(t)
            vx = np.zeros_like(t)
            vy = np.zeros_like(t)
            x = np.zeros_like(t)
            y = np.zeros_like(t)
        
        y = np.maximum(y, 0)
        
        # 火箭落地后停止模拟
        if np.any(y <= 0):
            landing_idx = np.argmax(y <= 0) + 1
            if landing_idx < len(t):
                t = t[:landing_idx]
                l = l[:landing_idx]
                vx = vx[:landing_idx]
                vy = vy[:landing_idx]
                x = x[:landing_idx]
                y = y[:landing_idx]
                
        # 确保数据完整
        min_len = min(len(t), len(l), len(x), len(y), len(vx), len(vy))
        t = t[:min_len]
        l = l[:min_len]
        x = x[:min_len]
        y = y[:min_len]
        vx = vx[:min_len]
        vy = vy[:min_len]
        
        # === 关键修复：计算加速度 ===
        acceleration = np.zeros_like(t)
        
        # 遍历所有时间点计算加速度
        for i in range(len(t)):
            # 重新计算当前状态的加速度
            l_i = l[i]
            vx_i = vx[i]
            vy_i = vy[i]
            x_i = x[i]
            y_i = y[i]
            
            # 火箭总质量
            water_mass = rho_water * S2 * l_i
            m = params['dry_mass'] + water_mass
            
            # 空气体积和压力
            V_air = max(1e-6, (l0 - l_i) * S2)
            p = pi * (Vi/V_air)**gamma
            
            # 水流条件
            water_flowing = (l_i > 1e-6) and (y_i > 0) and (p > p0)
            
            # 水流计算
            if water_flowing:
                pressure_diff = max(0, p - p0)
                v2 = np.sqrt(2 * pressure_diff / rho_water)
                dm_dt = -rho_water * S1 * v2
                thrust = -dm_dt * v2
            else:
                v2 = 0
                dm_dt = 0
                thrust = 0
            
            # 气动阻力
            v_speed = np.sqrt(vx_i**2 + vy_i**2)
            if v_speed > 0.01:
                eta = 0.5 * params['drag_coefficient'] * rho_air * S2
                Fr = eta * v_speed**2
                Fr_x = -Fr * (vx_i/v_speed)
                Fr_y = -Fr * (vy_i/v_speed)
            else:
                Fr_x, Fr_y = 0, 0
            
            # 推力方向
            if thrust > 0:
                if v_speed > 0.01:
                    thrust_x = thrust * (vx_i/v_speed)
                    thrust_y = thrust * (vy_i/v_speed)
                else:
                    thrust_x = thrust * np.cos(global_theta)
                    thrust_y = thrust * np.sin(global_theta)
            else:
                thrust_x, thrust_y = 0, 0
            
            # 加速度方程
            if y_i <= 0 and vy_i < 0:
                dvx_dt = 0
                dvy_dt = 0
            else:
                dvx_dt = (thrust_x + Fr_x) / m
                dvy_dt = (thrust_y - m*g + Fr_y) / m
            
            # 计算合加速度
            acceleration[i] = np.sqrt(dvx_dt**2 + dvy_dt**2)
    except Exception as e:
        print(f"模拟错误: {e}")
        t = np.array([0, 1])
        l = np.array([0, 0])
        x = np.array([0, 0])
        y = np.array([0, 0])
        vx = np.array([0, 0])
        vy = np.array([0, 0])
        acceleration = np.array([0, 0])
    
    # 模拟结束后确定每个时间点的阶段
    phases = []
    for i in range(len(t)):
        l_val = l[i]
        vy_val = vy[i]
        y_val = y[i]
        
        # 空气体积和压力
        V_air = max(1e-6, (l0 - l_val) * S2)
        p_val = pi * (Vi/V_air)**gamma
        
        # 判断当前阶段
        if y_val <= 0 and vy_val < 0:
            phases.append("着陆")
        elif l_val > 1e-6 and y_val > 0 and p_val > p0:
            phases.append("推进阶段")
        elif vy_val > 0:
            phases.append("惯性上升")
        else:
            phases.append("下降阶段")
    
    return t, x, y, vx, vy, phases, acceleration

# === 创建交互式图表 ===
fig = plt.figure(figsize=(16, 12), facecolor='white')

# 定义配色方案
primary_color = '#3498db'  # 主色调
secondary_color = '#2c3e50'  # 次级色调
text_color = '#2c3e50'  # 文本颜色
highlight_color = '#e74c3c'  # 高亮颜色

# 轨迹图区域 (顶部)
ax_traj = plt.axes([0.09, 0.55, 0.6, 0.35])
plt.title('水火箭仿真模拟器', fontsize=18, pad=15, color=text_color, fontweight='bold')
plt.xlabel('水平距离 (米)', fontsize=14, color=text_color)
plt.ylabel('高度 (米)', fontsize=14, color=text_color)
plt.grid(True, linestyle='--', alpha=0.3, color='#bdc3c7')

# 加速度图区域 (底部)
ax_accel = plt.axes([0.09, 0.10, 0.6, 0.35])
plt.title('加速度变化曲线', fontsize=16, pad=12, color=text_color, fontweight='bold')
plt.xlabel('时间 (秒)', fontsize=14, color=text_color, labelpad=10)
plt.ylabel('加速度 (m/s²)', fontsize=14, color=text_color, labelpad=10)
plt.grid(True, linestyle='--', alpha=0.3, color='#bdc3c7')

# 美化图表
for ax in [ax_traj, ax_accel]:
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_color('#7f8c8d')
    ax.spines['bottom'].set_color('#7f8c8d')
    ax.tick_params(axis='both', colors='#7f8c8d')
    ax.set_facecolor('white')
    ax.patch.set_edgecolor('#d6eaf8')
    ax.patch.set_linewidth(1.5)
    ax.patch.set_alpha(0.7)

ax_traj.set_xlim(0, 80)
ax_traj.set_ylim(0, 40)
ax_accel.set_xlim(0, 3)
ax_accel.set_ylim(0, 300)

# 按阶段绘制轨迹
phase_colors = {
    "推进阶段": '#3498db',  # 蓝色
    "惯性上升": '#2ecc71',  # 绿色
    "下降阶段": '#e74c3c',  # 红色
    "着陆": '#9b59b6'      # 紫色
}

# 初始模拟
t, x, y, vx, vy, phases, acceleration = simulate_rocket(params)

# 绘制轨迹
trajectory_lines = []
start_idx = 0
current_phase = phases[0]
for i in range(1, len(phases)):
    if phases[i] != current_phase:
        line, = ax_traj.plot(x[start_idx:i+1], y[start_idx:i+1], 
                       color=phase_colors[current_phase], linewidth=3,
                       path_effects=[path_effects.Stroke(linewidth=5, foreground='white'), 
                                     path_effects.Normal()])
        trajectory_lines.append(line)
        start_idx = i
        current_phase = phases[i]

if start_idx < len(phases):
    line, = ax_traj.plot(x[start_idx:], y[start_idx:], 
                   color=phase_colors[current_phase],
                   path_effects=[path_effects.Stroke(linewidth=5, foreground='white'), 
                                 path_effects.Normal()])
    trajectory_lines.append(line)

# 添加关键点
start_point, = ax_traj.plot([x[0]], [y[0]], 'g^', markersize=12, 
                     markeredgecolor='white', markeredgewidth=1.5)
end_point, = ax_traj.plot([x[-1]], [y[-1]], 'rs', markersize=12, 
                   markeredgecolor='white', markeredgewidth=1.5)
max_idx = np.argmax(y)
apogee_point, = ax_traj.plot([x[max_idx]], [y[max_idx]], highlight_color, marker='*', 
                      markersize=25, markeredgecolor='white',
                      path_effects=[path_effects.Stroke(linewidth=2, foreground='white'), 
                                    path_effects.Normal()])

# 添加地面
ground_level = 0
ax_traj.axhline(ground_level, color='#8B4513', linewidth=8, alpha=0.9,
           path_effects=[path_effects.Stroke(linewidth=9, foreground='#CD853F'), 
                         path_effects.Normal()])

# 添加发射台
ax_traj.plot([-1, 1], [ground_level, ground_level], '#2F4F4F', linewidth=5, alpha=0.9)
ax_traj.plot([0, 0], [ground_level, ground_level+0.2], '#708090', linewidth=3.5, alpha=0.9)
ax_traj.text(1.2, ground_level+0.1, '发射台', fontsize=12, color='#2F4F4F', fontweight='bold',
        path_effects=[path_effects.Stroke(linewidth=2, foreground='white'), 
                      path_effects.Normal()])

# === 添加关键参数信息框 ===
param_info_text = (
    f"水量比例: {params['water_fraction']*100:.1f}%\n"
    f"初始压力: {params['initial_pressure']*PSI_TO_ATM:.1f} PSI\n"
    f"发射角度: {params['launch_angle']:.1f}°\n"
    f"干质量: {params['dry_mass']:.3f} kg\n"
    f"瓶子容积: {params['bottle_volume']:.2f} L\n"
    f"瓶子直径: {params['bottle_diameter']:.3f} m\n"
    f"喷射口直径: {params['nozzle_diameter']:.4f} m\n"
    f"阻力系数: {params['drag_coefficient']:.2f}"
)

param_info_box = ax_traj.text(0.98, 0.95, param_info_text, transform=ax_traj.transAxes, 
                  fontsize=11, verticalalignment='top', 
                  horizontalalignment='right', color=text_color, fontweight='bold',
                  bbox=dict(boxstyle='round,pad=0.5', facecolor='white', 
                           edgecolor=primary_color, alpha=0.9, linewidth=2),
                  path_effects=[path_effects.Stroke(linewidth=1, foreground='white'), 
                                path_effects.Normal()])

# 添加加速度曲线
accel_line, = ax_accel.plot(t, acceleration, color='#e67e22', linewidth=2.5, 
                          path_effects=[path_effects.Stroke(linewidth=4, foreground='white'), 
                                        path_effects.Normal()])
max_accel_point, = ax_accel.plot([t[np.argmax(acceleration)]], [np.max(acceleration)], 
                               color='#c0392b', marker='o', markersize=10,
                               markeredgecolor='white', markeredgewidth=1.5)

# 添加加速度信息
if len(acceleration) > 0:
    max_accel = np.max(acceleration)
    mean_accel = np.mean(acceleration)
    accel_info = (
        f"最大加速度: {max_accel:.1f} m/s² ({max_accel/g:.1f} g)\n"
        f"平均加速度: {mean_accel:.1f} m/s²"
    )
else:
    accel_info = "无加速度数据"

accel_info_box = ax_accel.text(0.98, 0.95, accel_info, transform=ax_accel.transAxes, 
                             fontsize=12, verticalalignment='top', 
                             horizontalalignment='right', color=text_color, fontweight='bold',
                             bbox=dict(boxstyle='round,pad=0.5', facecolor='white', 
                                      edgecolor='#e67e22', alpha=0.9, linewidth=2),
                             path_effects=[path_effects.Stroke(linewidth=1, foreground='white'), 
                                           path_effects.Normal()])

# === 右侧垂直控件区域 ===
control_panel_left = 0.71
control_panel_top = 0.90
slider_height = 0.03
slider_width = 0.25
slider_spacing = 0.065
slider_bg_color = '#ecf0f1'
slider_title_fontsize = 11

# 创建滑块框架
slider_frame = plt.axes([control_panel_left - 0.04, 0.10, 0.30, 0.78], facecolor='none')
slider_frame.set_axis_off()

frame_patch = FancyBboxPatch(
    (control_panel_left - 0.04, 0.10), 0.30, 0.78,
    boxstyle="round,pad=0.02", facecolor='#f8f9f9',
    edgecolor='#d6eaf8', alpha=0.9, linewidth=2
)
slider_frame.add_patch(frame_patch)

slider_frame.text(control_panel_left + 0.135, 0.87, '控制面板', 
                 fontsize=16, color=text_color, fontweight='bold',
                 ha='center', va='center',
                 bbox=dict(boxstyle='round,pad=0.5', facecolor=primary_color, 
                          edgecolor='white', alpha=0.9, linewidth=0.8),
                 path_effects=[path_effects.Stroke(linewidth=1, foreground='white'), 
                               path_effects.Normal()])

# 创建滑块函数
def create_slider_with_label(position, title, valmin, valmax, valinit, valfmt=None):
    ax_slider = plt.axes([control_panel_left, position, slider_width, slider_height], 
                         facecolor=slider_bg_color)
    slider = Slider(ax_slider, '', valmin, valmax, valinit=valinit, 
                   valfmt=valfmt, color=primary_color, track_color='#d6eaf8')
    ax_slider.text(0.5, 0, title, transform=ax_slider.transAxes, 
                  ha='center', va='top', fontsize=slider_title_fontsize, 
                  color=text_color, fontweight='bold')
    return slider

# 水量滑块
water_slider = create_slider_with_label(
    control_panel_top, 
    '水量比例 (%)', 
    10, 80, 
    params['water_fraction']*100, 
    '%.0f %%'
)

# 压力滑块（使用PSI为单位）
min_pressure_psi = 30
max_pressure_psi = 90
initial_pressure_psi = params['initial_pressure'] * PSI_TO_ATM
pressure_slider = create_slider_with_label(
    control_panel_top - slider_spacing, 
    '初始压力 (PSI)', 
    min_pressure_psi, max_pressure_psi, 
    initial_pressure_psi
)

# 角度滑块
angle_slider = create_slider_with_label(
    control_panel_top - 2*slider_spacing, 
    '发射角度 (°)', 
    0, 90, 
    params['launch_angle']
)

# 质量滑块
mass_slider = create_slider_with_label(
    control_panel_top - 3*slider_spacing, 
    '干质量 (kg)', 
    0.01, 0.5, 
    params['dry_mass']
)

# 容积滑块
volume_slider = create_slider_with_label(
    control_panel_top - 4*slider_spacing, 
    '瓶子容积 (L)', 
    0.5, 2.5, 
    params['bottle_volume']
)

# 直径滑块
diameter_slider = create_slider_with_label(
    control_panel_top - 5*slider_spacing, 
    '瓶子直径 (m)', 
    0.05, 0.15, 
    params['bottle_diameter']
)

# 喷射口直径滑块
nozzle_slider = create_slider_with_label(
    control_panel_top - 6*slider_spacing, 
    '喷射口直径 (m)', 
    0.005, 0.03, 
    params['nozzle_diameter']
)

# 阻力滑块
drag_slider = create_slider_with_label(
    control_panel_top - 7*slider_spacing, 
    '阻力系数', 
    0.1, 1.0,params['drag_coefficient']
)

# 重置按钮
reset_ax = plt.axes([control_panel_left + 0.05, control_panel_top - 8.5*slider_spacing, 
                     0.15, 0.05])
reset_button = Button(reset_ax, '重置参数', color=primary_color, 
                     hovercolor=highlight_color)
reset_button.label.set_color('white')
reset_button.label.set_fontweight('bold')
reset_button.label.set_fontsize(12)
reset_ax.patch.set_path_effects([path_effects.Stroke(linewidth=1, foreground='white'), 
                                path_effects.Normal()])

# === 参数分析按钮 ===
analysis_ax = plt.axes([control_panel_left + 0.05, control_panel_top - 9.5*slider_spacing, 
                        0.15, 0.05])
analysis_button = Button(analysis_ax, '参数分析', color='#27ae60', 
                     hovercolor='#2ecc71')
analysis_button.label.set_color('white')
analysis_button.label.set_fontweight('bold')
analysis_button.label.set_fontsize(12)
analysis_ax.patch.set_path_effects([path_effects.Stroke(linewidth=1, foreground='white'), 
                                   path_effects.Normal()])

# === 最大加速度分析按钮 ===
max_accel_ax = plt.axes([control_panel_left + 0.05, control_panel_top - 10.5*slider_spacing, 
                        0.15, 0.05])
max_accel_button = Button(max_accel_ax, '最大加速度分析', color='#9b59b6', 
                     hovercolor='#8e44ad')
max_accel_button.label.set_color('white')
max_accel_button.label.set_fontweight('bold')
max_accel_button.label.set_fontsize(12)
max_accel_ax.patch.set_path_effects([path_effects.Stroke(linewidth=1, foreground='white'), 
                                   path_effects.Normal()])

# 创建大气压说明文本
pressure_info_text = plt.figtext(0.72, 0.05, 
            f'压强单位说明: 1 ATM = {PSI_TO_ATM:.1f} PSI\n当前压力: {initial_pressure_psi:.1f} PSI ≈ {params["initial_pressure"]:.1f} ATM', 
            fontsize=10, color=text_color, ha='center')

# === 修复图例重叠问题 ===
# 添加阶段图例 - 移动到左上角
phase_legend = ax_traj.legend(handles=[
    Line2D([0], [0], color=phase_colors["推进阶段"], lw=3, label='推进阶段'),
    Line2D([0], [0], color=phase_colors["惯性上升"], lw=3, label='惯性上升'),
    Line2D([0], [0], color=phase_colors["下降阶段"], lw=3, label='下降阶段'),
    Line2D([0], [0], color=phase_colors["着陆"], lw=3, label='着陆')
], loc='upper left', fontsize=10, framealpha=0.95, 
           facecolor='white', edgecolor=text_color,
           bbox_to_anchor=(0.02, 0.98))

# 添加最高点图例 - 移动到右上角
apogee_legend = ax_traj.legend(handles=[
    Line2D([0], [0], marker='*', color=highlight_color, lw=0, markersize=25, 
           markeredgewidth=1.5, markeredgecolor='white', label='最高点')
], loc='upper right', fontsize=10, framealpha=0.95, 
           facecolor='white', edgecolor=text_color,
           bbox_to_anchor=(0.98, 0.98))

ax_traj.add_artist(phase_legend)
ax_traj.add_artist(apogee_legend)

# 更新函数
def update(val):
    # 更新参数
    params['water_fraction'] = water_slider.val / 100
    params['initial_pressure'] = pressure_slider.val / PSI_TO_ATM
    params['launch_angle'] = angle_slider.val
    params['dry_mass'] = mass_slider.val
    params['bottle_volume'] = volume_slider.val
    params['bottle_diameter'] = diameter_slider.val
    params['nozzle_diameter'] = nozzle_slider.val
    params['drag_coefficient'] = drag_slider.val
    
    # 使用新参数模拟
    t, x, y, vx, vy, phases, acceleration = simulate_rocket(params)
    
    # 清除旧轨迹
    for line in trajectory_lines:
        line.remove()
    trajectory_lines.clear()
    
    # 按阶段绘制新轨迹
    start_idx = 0
    current_phase = phases[0]
    for i in range(1, len(phases)):
        if phases[i] != current_phase:
            line, = ax_traj.plot(x[start_idx:i+1], y[start_idx:i+1], 
                           color=phase_colors[current_phase], linewidth=3,
                           path_effects=[path_effects.Stroke(linewidth=5, foreground='white'), 
                                        path_effects.Normal()])
            trajectory_lines.append(line)
            start_idx = i
            current_phase = phases[i]
    
    if start_idx < len(phases):
        line, = ax_traj.plot(x[start_idx:], y[start_idx:], 
                       color=phase_colors[current_phase], linewidth=3,
                       path_effects=[path_effects.Stroke(linewidth=5, foreground='white'), 
                                    path_effects.Normal()])
        trajectory_lines.append(line)
    
    # 更新关键点
    start_point.set_data([x[0]], [y[0]])
    end_point.set_data([x[-1]], [y[-1]])
    
    if len(y) > 0:
        max_idx = np.argmax(y)
        apogee_point.set_data([x[max_idx]], [y[max_idx]])
    else:
        apogee_point.set_data([], [])
    
    # 更新参数信息框
    param_info_text = (
        f"水量比例: {params['water_fraction']*100:.1f}%\n"
        f"初始压力: {params['initial_pressure']*PSI_TO_ATM:.1f} PSI\n"
        f"发射角度: {params['launch_angle']:.1f}°\n"
        f"干质量: {params['dry_mass']:.3f} kg\n"
        f"瓶子容积: {params['bottle_volume']:.2f} L\n"
        f"瓶子直径: {params['bottle_diameter']:.3f} m\n"
        f"喷射口直径: {params['nozzle_diameter']:.4f} m\n"
        f"阻力系数: {params['drag_coefficient']:.2f}"
    )
    param_info_box.set_text(param_info_text)
    
    # 更新加速度曲线
    if len(t) > 0 and len(acceleration) > 0:
        accel_line.set_data(t, acceleration)
        if len(acceleration) > 0:
            max_idx = np.argmax(acceleration)
            max_accel_point.set_data([t[max_idx]], [acceleration[max_idx]])
        else:
            max_accel_point.set_data([], [])
        
        # 更新加速度信息
        if len(acceleration) > 0:
            max_accel = np.max(acceleration)
            mean_accel = np.mean(acceleration)
            accel_info = (
                f"最大加速度: {max_accel:.1f} m/s² ({max_accel/g:.1f} g)\n"
                f"平均加速度: {mean_accel:.1f} m/s²"
            )
        else:
            accel_info = "无加速度数据"
        
        accel_info_box.set_text(accel_info)
        
        # 计算飞行时间并确定合适的显示范围
        flight_time = t[-1] if len(t) > 0 else 0
        
        # 动态调整时间范围
        if flight_time <= 3:
            time_limit = max(1.0, flight_time * 1.2)
        else:
            time_limit = min(15.0, flight_time * 1.05)
        
        ax_accel.set_xlim(0, time_limit)
        
        # 设置y轴范围
        if len(acceleration) > 0:
            y_max = max(300, np.max(acceleration)*1.2)
            ax_accel.set_ylim(0, y_max)
        else:
            ax_accel.set_ylim(0, 300)
    else:
        accel_line.set_data([], [])
        max_accel_point.set_data([], [])
        accel_info_box.set_text("无加速度数据")
        ax_accel.set_xlim(0, 3)
        ax_accel.set_ylim(0, 300)
    
    # 更新大气压说明文本
    pressure_info_text.set_text(
        f'压强单位说明: 1 ATM = {PSI_TO_ATM:.1f} PSI\n当前压力: {pressure_slider.val:.1f} PSI ≈ {params["initial_pressure"]:.1f} ATM'
    )
    
    # 调整轨迹图坐标轴范围
    if len(x) > 0 and len(y) > 0:
        margin_x = max(5, max(x) * 0.2)  # 20%的边距
        margin_y = max(5, max(y) * 0.2)  # 20%的边距
        ax_traj.set_xlim(0, max(x) + margin_x)
        ax_traj.set_ylim(0, max(y) + margin_y)
    else:
        # 如果没有有效数据，重置为初始范围
        ax_traj.set_xlim(0, 80)
        ax_traj.set_ylim(0, 40)
    
    fig.canvas.draw_idle()

# === 确保滑块连接到更新函数 ===
water_slider.on_changed(update)
pressure_slider.on_changed(update)
angle_slider.on_changed(update)
mass_slider.on_changed(update)
volume_slider.on_changed(update)
diameter_slider.on_changed(update)
nozzle_slider.on_changed(update)
drag_slider.on_changed(update)

# === 优化后的参数影响分析函数 ===
def plot_parameter_effects(event):
    # 创建一个新的窗口（figure）使用白色背景
    fig2 = plt.figure(figsize=(16, 16), facecolor='white')
    fig2.canvas.manager.set_window_title('水火箭参数影响分析')
    
    # 使用GridSpec创建布局 - 大幅增加垂直间距
    gs = GridSpec(3, 4, figure=fig2, hspace=0.8, wspace=0.3)
    
    # 添加主标题
    title_ax = fig2.add_subplot(gs[0, :])
    title_ax.set_axis_off()
    title_text = title_ax.text(0.5, 0.5, '水火箭参数影响分析', 
                              fontsize=24, ha='center', va='center',
                              fontweight='bold', color='#333333')
    title_text.set_path_effects([path_effects.Stroke(linewidth=3, foreground='white'),
                               path_effects.Normal()])
    
    # 定义需要分析的参数及其值范围
    param_info = [
        ('water_fraction', '水量比例', '%', [0.2, 0.3, 0.4, 0.5, 0.6], lambda x: x, 'upper right'),
        ('initial_pressure', '初始压力', 'ATM', [2.0, 3.0, 4.0, 5.0, 6.0], lambda x: x, 'upper right'),
        ('launch_angle', '发射角度', '°', [30, 40, 50, 60, 70], lambda x: x, 'upper right'),
        ('dry_mass', '干质量', 'kg', [0.03, 0.06, 0.09, 0.12, 0.15], lambda x: x, 'upper right'),
        ('bottle_volume', '瓶子容积', 'L', [0.25, 0.5, 0.75, 1.0, 1.25], lambda x: x, 'upper right'),
        ('bottle_diameter', '瓶子直径', 'm', [0.07, 0.09, 0.11, 0.13, 0.15], lambda x: x, 'upper right'),
        ('nozzle_diameter', '喷射口直径', 'm', [0.007, 0.011, 0.015, 0.019, 0.023], lambda x: x, 'upper right'),
        ('drag_coefficient', '阻力系数', '', [0.2, 0.4, 0.6, 0.8, 1.0], lambda x: x, 'upper right')
    ]
    
    # 创建8个子图 - 修复索引问题
    plot_axes = []
    for i in range(8):
        row = i // 4 + 1  # 计算行索引 (1 或 2)
        col = i % 4       # 计算列索引 (0-3)
        ax = fig2.add_subplot(gs[row, col])
        ax.set_facecolor('white')  # 设置子图背景为白色
        plot_axes.append(ax)
    
    # 为每个参数生成加速度曲线
    for ax, (param_name, param_label, unit, param_values, conv_func, legend_loc) in zip(plot_axes, param_info):
        # 设置专业样式的子图
        ax.set_facecolor('#fafafa')  # 非常浅的灰色背景
        ax.grid(True, linestyle='--', alpha=0.7, color='#d0d0d0')  # 浅灰色网格
        
        # 使用科学期刊风格的颜色
        colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd']
        line_styles = ['-', '--', '-.', ':', '-']
        
        # 记录该子图的最大加速度和推进阶段持续时间
        max_accel_in_ax = 0
        max_thrust_duration = 0
        
        for i, value in enumerate(param_values):
            # 使用当前主窗口参数作为基础
            test_params = params.copy()
            test_params[param_name] = conv_func(value)
            
            try:
                # 进行模拟
                t, x, y, vx, vy, phases, acceleration = simulate_rocket(test_params)
                
                # 忽略无效数据
                if len(t) < 2 or len(acceleration) < 2:
                    continue
                    
                # 记录最大加速度
                current_max_accel = np.max(acceleration)
                if current_max_accel > max_accel_in_ax:
                    max_accel_in_ax = current_max_accel
                
                # 计算推进阶段持续时间
                thrust_duration = 0
                for j in range(len(phases)):
                    if phases[j] == "推进阶段":
                        thrust_duration = t[j]
                        break
                if thrust_duration > max_thrust_duration:
                    max_thrust_duration = thrust_duration
                
                # 创建平滑曲线（使用样条插值）
                if len(t) > 10:
                    # 创建时间轴插值点 - 只关注推进阶段
                    new_t = np.linspace(0, min(1.0, t[-1]), 200)  # 限制在1秒内
                    spl = make_interp_spline(t, acceleration, k=3)
                    accel_smooth = spl(new_t)
                    
                    # 限制加速度不小于0
                    accel_smooth = np.maximum(accel_smooth, 0)
                    
                    # 绘制平滑曲线
                    line, = ax.plot(new_t, accel_smooth, color=colors[i % len(colors)], 
                                    linestyle=line_styles[i % len(line_styles)],
                                    linewidth=2.5, alpha=0.9, 
                                    label=f'{value} {unit}')
                    
                    # 标记最大加速度点（使用原始数据）
                    max_idx = np.argmax(acceleration)
                    if t[max_idx] <= new_t[-1]:  # 确保点在可视范围内
                        ax.plot(t[max_idx], acceleration[max_idx], 
                                'o', markersize=8, color=colors[i % len(colors)],
                                markeredgecolor='white', markeredgewidth=1)
                
            except Exception as e:
                print(f"参数分析错误: {param_name}={value}, {e}")
                continue
        
        # 设置专业样式的标题 - 增加标题与子图的间距
        ax.set_title(f'{param_label}的影响', fontsize=14, pad=15, 
                    color='#333333', fontweight='bold')
        
        # 设置专业样式的坐标轴标签
        ax.set_xlabel('时间 (秒)', fontsize=12, color='#555555', labelpad=10)
        ax.set_ylabel('加速度 (m/s²)', fontsize=12, color='#555555', labelpad=12)
        
        # 设置坐标轴范围 - 专注于推进阶段
        # 时间轴：0到推进阶段持续时间+0.1秒，最大不超过1秒
        time_limit = min(1.0, max_thrust_duration + 0.1) if max_thrust_duration > 0 else 0.5
        ax.set_xlim(0, time_limit)
        
        # Y轴范围：0到最大加速度的1.2倍
        if max_accel_in_ax > 0:
            ax.set_ylim(0, max_accel_in_ax * 1.2)
        else:
            ax.set_ylim(0, 300)
        
        # 添加图例 - 使用更专业的样式
        if ax.get_legend_handles_labels()[0]:
            legend = ax.legend(loc=legend_loc, fontsize=9, framealpha=0.95,
                             facecolor='white', edgecolor='#cccccc', frameon=True,
                             title=f'{param_label} ({unit})', title_fontsize=11)
            legend.get_title().set_fontweight('bold')
            
            # 调整图例位置避免遮挡曲线
            if legend_loc == 'upper right':
                legend.set_bbox_to_anchor((1.02, 1.0))  # 稍微向右移动
    
    # 添加页脚说明
    fig2.text(0.5, 0.02, 
             '水火箭参数影响分析 | 专注于推进阶段 (0-1秒) | 点标记为最大加速度点',
             ha='center', fontsize=12, color='#777777', alpha=0.8)
    
    # 优化布局 - 增加顶部空间避免标题重叠
    plt.tight_layout(rect=[0, 0.03, 1, 0.95])
    
    # 显示新窗口
    fig2.canvas.draw_idle()
    plt.show()

# === 优化后的最大加速度分析函数 ===
def plot_max_acceleration_trend(event):
    # 创建一个新的窗口（figure）
    fig3 = plt.figure(figsize=(16, 20), facecolor='white')
    fig3.canvas.manager.set_window_title('水火箭最大加速度参数分析')
    
    # 使用GridSpec创建2x4布局 - 大幅增加垂直间距
    gs = GridSpec(2, 4, figure=fig3, hspace=0.8, wspace=0.3)
    
    # 添加主标题
    title_ax = fig3.add_subplot(gs[0, :])
    title_ax.set_axis_off()
    title_text = title_ax.text(0.5, 0.5, '水火箭参数对最大加速度的影响', 
                              fontsize=24, ha='center', va='center',
                              fontweight='bold', color='#2c3e50')
    title_text.set_path_effects([path_effects.Stroke(linewidth=3, foreground='white'),
                               path_effects.Normal()])
    
    # 定义关键参数
    key_params = [
        ('water_fraction', '水量比例', '%', 1, 80, lambda x: x/100),
        ('initial_pressure', '初始压力', 'PSI', 30, 90, lambda x: x/PSI_TO_ATM),
        ('dry_mass', '干质量', 'kg', 0.01, 0.5, lambda x: x),
        ('bottle_volume', '瓶子容积', 'L', 0.5, 2.5, lambda x: x),
        ('bottle_diameter', '瓶子直径', 'm', 0.05, 0.15, lambda x: x),
        ('nozzle_diameter', '喷射口直径', 'm', 0.007, 0.025, lambda x: x),
        ('drag_coefficient', '阻力系数', '', 0.1, 1.0, lambda x: x),
        ('launch_angle', '发射角度', '°', 0, 90, lambda x: x)
    ]
    
    # 创建8个子图
    plot_axes = [fig3.add_subplot(gs[1, i]) for i in range(4)] + [fig3.add_subplot(gs[2, i]) for i in range(4)]
    
    # 为每个参数生成曲线
    for ax, (param_name, param_label, unit, min_val, max_val, conv_func) in zip(plot_axes, key_params):
        try:
            # 增加采样点
            param_values = np.linspace(min_val, max_val, 15)
            
            # 存储最大加速度值
            max_accels = []
            
            for value in param_values:
                test_params = params.copy()
                test_params[param_name] = conv_func(value)
                
                try:
                    # 进行模拟
                    t, x, y, vx, vy, phases, acceleration = simulate_rocket(test_params)
                    
                    # 找出最大加速度
                    if len(acceleration) > 0:
                        max_accel = np.max(acceleration)
                        max_accels.append(max_accel)
                    else:
                        max_accels.append(0)
                        
                except:
                    max_accels.append(0)
            
            # 设置专业样式
            ax.set_facecolor('#fafafa')
            ax.grid(True, linestyle='--', alpha=0.7, color='#d0d0d0')
            
            # 使用科学期刊风格的颜色
            ax.plot(param_values, max_accels, 'o-', color='#1f77b4', linewidth=2, markersize=5)
            
            # 设置标题和标签 - 增加标题与子图的间距
            ax.set_title(f'{param_label}对最大加速度的影响', fontsize=14, pad=15, 
                        color='#333333', fontweight='bold')
            ax.set_xlabel(f'{param_label} ({unit})', fontsize=12, color='#555555', labelpad=10)
            ax.set_ylabel('最大加速度 (m/s²)', fontsize=12, color='#555555', labelpad=10)
            
            # 设置y轴范围
            if max(max_accels) > 0:
                ax.set_ylim(0, max(max_accels)*1.2)
            
        except Exception as e:
            print(f"绘制{param_label}曲线时出错: {e}")
            continue
    
    # 添加页脚说明
    fig3.text(0.5, 0.02, 
             '水火箭参数对最大加速度影响分析 | 基于主窗口当前设置',
             ha='center', fontsize=12, color='#777777', alpha=0.9)
    
    # 优化布局
    plt.tight_layout(rect=[0, 0.03, 1, 0.95])
    
    # 显示新窗口
    fig3.canvas.draw_idle()
    plt.show()

# === 重置函数 ===
def reset(event):
    # 重置参数为初始值
    params.update(initial_params)
    
    # 重置滑块而不触发更新
    sliders = [water_slider, pressure_slider, angle_slider, 
               mass_slider, volume_slider, diameter_slider,
               nozzle_slider, drag_slider]
    
    for slider in sliders:
        slider.eventson = False
    
    # 设置滑块值（PSI滑块需要转换）
    water_slider.set_val(initial_params['water_fraction'] * 100)
    # 将ATM转换为PSI（使用转换系数）
    pressure_slider.set_val(initial_params['initial_pressure'] * PSI_TO_ATM)
    angle_slider.set_val(initial_params['launch_angle'])
    mass_slider.set_val(initial_params['dry_mass'])
    volume_slider.set_val(initial_params['bottle_volume'])
    diameter_slider.set_val(initial_params['bottle_diameter'])
    nozzle_slider.set_val(initial_params['nozzle_diameter'])
    drag_slider.set_val(initial_params['drag_coefficient'])
    
    for slider in sliders:
        slider.eventson = True
    
    # 手动更新到初始状态
    update(None)

# 连接按钮到函数
reset_button.on_clicked(reset)
analysis_button.on_clicked(plot_parameter_effects)
max_accel_button.on_clicked(plot_max_acceleration_trend)

# 显示图表
plt.show()